#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys

import pygtk
pygtk.require('2.0')
import gtk
import gobject

import hashlib
import shutil
import tarfile

import subprocess

#Store = { 'path': "", 'hash': "" }
#View = { path, ext, size, hash }

'''
Data Store:
    All the data saved in a database, like git. All the files are hashed, the real
    file just link to that hashed file.

    Every hash file as a ref count
'''

if os.path.islink(__file__):
    this_file = os.readlink(__file__)
else:
    this_file = __file__

def md5_file(f, block_size=2**20):
    md5 = hashlib.md5()
    while True:
        data = f.read(block_size)
        if not data:
            break
        md5.update(data)
    return md5.hexdigest()

def tar_folder(folder, target):
    tf = tarfile.open(target, 'w:bz2')
    tf.add(folder)
    tf.close()

class ProjectContent:
    def __init__(self, path):
        self.__files = {}
        self.__path = path
        self.__index_file = os.path.join(path, 'bf.index')
        self.__blobs_dir = os.path.join(path, 'bf.blobs')

    def get_index_file(self):
        return self.__index_file

    def get_blobs_dir(self):
        return self.__blobs_dir

    def get_path(self):
        return self.__path

    def get_files(self):
        return self.__files

    def add_path(self, path, hash = '0' * 32):
        if path != '':
            self.__files[path.strip()] = hash.strip()

    def del_path(self, path):
        del self.__files[path]

class MainWindow(gtk.Window):
    __project_content = None

    # changed against the original project
    __dirty = 3

    def load(self):
        files = {}

        try:
            f = open(self.__project_content.get_index_file(), 'a+')
        except e:
            print e

        lines = f.readlines()
        for line in lines:
            ''' line = hash + ' ' + filename '''
            if len(line) < 33:
                continue

            hash = line[:32]
            path = line[33:]

            if path != '':
                files[path.strip()] = hash.strip()
        f.close()
        return files

    def save(self):
        try:
            f = open(self.__project_content.get_index_file(), 'w+')
        except e:
            return

        for k, v in self.__project_content.get_files().items():
            line = '%s %s\n' % (v, k)
            f.writelines([line])

        f.close()

    def __set_dirty(self, set):
        if set == self.__dirty:
            return

        if set:
            self.__dirty = True
        else:
            self.__dirty = False

        self.tool_bar_save_project.set_sensitive(self.__dirty)

    def __ena(self, button):
        button.set_sensitive(True)

    def __dis(self, button):
        button.set_sensitive(False)

    def __on_quit(self, win, evt = gtk.gdk.DELETE):
        gtk.main_quit(win)

    def __change_title(self, title='Bei the shit'):
        ''' Change the window title to current project path '''
        self.set_title(title)

    def __init__(self):
        gtk.Window.__init__(self)
        self.connect('delete_event', self.__on_quit)

        self.__change_title()

        self.set_default_size(800, 480)
        self.set_size_request(800, 480)

        self.__body_vbox = gtk.VBox()
        self.add(self.__body_vbox)

        self.__create_tool_bar()
        self.__body_vbox.pack_start(gtk.HSeparator(), expand=False, padding=4)
        self.__create_proj_list()
        self.__create_status_bar()

        settings = self.get_settings()
        settings.set_property("gtk-button-images", True)

        self.show_all()

    def __load_files_to_view(self):
        files = self.load()
        for f, h in files.items():
            # add_directory will cause add recursive, DO NOT use
            self.add_files([f])

    def __ask_close_project(self):
        dlg = gtk.MessageDialog(None, gtk.DIALOG_MODAL, gtk.MESSAGE_WARNING, gtk.BUTTONS_YES_NO)
        dlg.set_markup('<big><b>%s</b></big>' % 'Warning')
        dlg.format_secondary_markup('Close project %s ?' % self.__project_content.get_path())

        ret = False
        if dlg.run() == gtk.RESPONSE_YES:
            ret = True
        dlg.destroy()
        return ret

    def __close_project(self):
        if self.__ask_close_project():
            self.del_all()
            self.__project_content = None
            self.__change_title()

            self.__ena(self.tool_bar_new_project)
            self.__dis(self.tool_bar_close_project)
            self.__dis(self.tool_bar_save_project)
            self.__dis(self.tool_bar_backup_project)
            self.__dis(self.tool_bar_search_project)
            self.__dis(self.tool_bar_check_project)
            self.__dis(self.tool_bar_add_file)
            self.__dis(self.tool_bar_add_directory)
            self.__dis(self.tool_bar_expand_all)
            self.__dis(self.tool_bar_collapse_all)
            self.__dis(self.tool_bar_load_cache)

    def __on_new_project(self, button):
        ''' Open or create a project_content, should close previous first '''

        if self.__project_content:
            if self.__ask_close_project():
                self.__close_project()
                self.__project_content = None
            else:
                return

        file_dialog = gtk.FileChooserDialog(
            title = "Create or open project",
            action = gtk.FILE_CHOOSER_ACTION_CREATE_FOLDER,
            buttons = (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))

        if file_dialog.run() == gtk.RESPONSE_OK:
            path = file_dialog.get_filename()

            self.__proj_path = path

            self.__project_content = ProjectContent(path)
            self.__load_files_to_view()

            self.__change_title(path)

            self.__dis(self.tool_bar_new_project)
            self.__ena(self.tool_bar_close_project)
            self.__dis(self.tool_bar_save_project)
            self.__ena(self.tool_bar_backup_project)
            self.__ena(self.tool_bar_search_project)
            self.__ena(self.tool_bar_check_project)
            self.__ena(self.tool_bar_add_file)
            self.__ena(self.tool_bar_add_directory)
            self.__ena(self.tool_bar_expand_all)
            self.__ena(self.tool_bar_collapse_all)
            self.__ena(self.tool_bar_load_cache)

            self.__set_dirty(False)

        file_dialog.destroy()

    def __on_close_project(self, button):
        if self.__project_content:
            self.__close_project()
            self.__proj_path = None

    def __on_save_project(self, button):
        if self.__project_content:
            self.save()
            self.__set_dirty(False)

    def __backup_project_to_file(self, target_fn):
        if not target_fn.endswith('.tar.bz2'):
            target_fn = target_fn + '.tar.bz2'

        restore_py = os.path.join(os.path.realpath(os.path.dirname(__file__)), 'restore.py')

        blobdir = self.__project_content.get_blobs_dir()

        # empty the target directory
        if os.path.exists(blobdir):
            if os.path.isdir(blobdir):
                shutil.rmtree(blobdir)
            else:
                os.remove(blobdir)

        os.mkdir(blobdir)

        self.__dis(self.tool_bar_backup_project)

        for path, v in self.__project_content.get_files().items():
            try:
                fh = open(path, 'r')
            except:
                continue
            hash = md5_file(fh)
            fh.close()

            self.__project_content.add_path(path, hash)

            target_blob = os.path.join(blobdir, hash)
            if not os.path.exists(target_blob):
                self.__set_status_text("Copy '%s' to '%s'" % (path, target_blob))
                shutil.copy(path, target_blob)
            else:
                self.__set_status_text("'%s' already exist" % target_blob)

        self.save()
        self.__set_dirty(False)

        self.__ena(self.tool_bar_backup_project)

        self.__set_status_text("Zipping target '%s'" % target_fn)

        curdir = os.curdir
        os.chdir(os.path.dirname(self.__proj_path))
        shutil.copy(restore_py, blobdir + '/..')
        tar_folder(os.path.basename(self.__proj_path), target_fn)
        os.remove(os.path.join(blobdir + '/..', 'restore.py'))
        os.chdir(curdir)

    def __on_backup_project(self, button):
        if not self.__project_content:
            return

        if self.__dirty:
            dlg = gtk.MessageDialog(None, gtk.DIALOG_MODAL, gtk.MESSAGE_WARNING, gtk.BUTTONS_OK)
            dlg.set_markup('<big><b>%s</b></big>' % 'Warning')
            dlg.format_secondary_markup('Should save project first.')
            dlg.run()
            dlg.destroy()
            return

        file_dialog = gtk.FileChooserDialog(
            title = 'Select the target zip file',
            action = gtk.FILE_CHOOSER_ACTION_SAVE,
            buttons = (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))

        if file_dialog.run() == gtk.RESPONSE_OK:
            target_fn = file_dialog.get_filename()
            file_dialog.destroy()

            self.__backup_project_to_file(target_fn)

            dlg = gtk.MessageDialog(None, gtk.DIALOG_MODAL, gtk.MESSAGE_WARNING, gtk.BUTTONS_OK)
            dlg.set_markup('<big><b>%s</b></big>' % 'Backup complete')
            dlg.format_secondary_markup('From [%s] to [%s].' % (self.__proj_path, target_fn))
            dlg.run()
            dlg.destroy()

        else:
            file_dialog.destroy()

    def __on_search_project(self, button):
        ''' Show input entry dialog, the find the files by nm: or szgt: or szlt: '''
        ''' Hilight the matched lines '''
        if not self.__project_content:
            return

        dialog = SearchWindow()
        if dialog.run() == gtk.RESPONSE_OK:
            self.search(dialog.get_text())

        dialog.destroy()

    def __on_check_project(self, button):
        if not self.__project_content:
            return

        def check_exist(model, path, iter, selection):
            path = model.get(iter, 1)[0]
            if not os.path.exists(path):
                selection.select_iter(iter)

        selection = self.__proj_tree_view.get_selection()
        selection.unselect_all()
        self.__proj_tree_view.get_model().foreach(check_exist, selection)

    def __on_add_files(self, button):
        if not self.__project_content:
            return

        file_dialog = gtk.FileChooserDialog(
            title = "Add files",
            action = gtk.FILE_CHOOSER_ACTION_OPEN,
            buttons = (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))
        file_dialog.set_select_multiple(True)

        if file_dialog.run() == gtk.RESPONSE_OK:
            self.add_files(file_dialog.get_filenames())

        file_dialog.destroy()

    def __on_add_directory(self, button):
        if not self.__project_content:
            return

        file_dialog = gtk.FileChooserDialog(
            title = "Add directories",
            action = gtk.FILE_CHOOSER_ACTION_SELECT_FOLDER,
            buttons = (gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL, gtk.STOCK_OPEN, gtk.RESPONSE_OK))
        file_dialog.set_select_multiple(True)

        if file_dialog.run() == gtk.RESPONSE_OK:
            self.add_directory(file_dialog.get_filenames()[0])

        file_dialog.destroy()

    def __on_expand_all(self, button):
        self.__proj_tree_view.expand_all()

    def __on_collapse_all(self, button):
        self.__proj_tree_view.collapse_all()

    def __on_load_cache(self, button):

        try:
            f = open(os.path.expanduser('~/.beifen-cache'), 'a+')
        except:
            return

        files = []
        lines = f.readlines()
        f.close()

        for line in lines:
            line = line.strip()

            if len(line) < 1 or not os.path.exists(line):
                continue

            files.append(line)

        file_count = len(files)
        if file_count <= 0:
            return

        # Message for user to continue
        dlg = gtk.MessageDialog(None, gtk.DIALOG_MODAL, gtk.MESSAGE_WARNING, gtk.BUTTONS_YES_NO)
        dlg.set_markup('<big><b>%d files or/and directories found in cache</b></big>' % file_count)
        dlg.format_secondary_markup('Continue to load?')

        if dlg.run() == gtk.RESPONSE_YES:
            for f in files:
                if os.path.isdir(f):
                    self.add_directory(f)
                else:
                    self.add_files([f])

        dlg.destroy()

    def __create_tool_bar(self):
        tool_bar_hbox = gtk.HBox()
        tool_bar_hbox.set_size_request(-1, 36)

        def stock(stockname):
            img = gtk.Image()
            img.set_from_stock(stockname, gtk.ICON_SIZE_BUTTON)
            return img

        self.tool_bar_new_project = gtk.Button('_Open / New')
        self.tool_bar_new_project.set_image(stock(gtk.STOCK_OPEN))
        self.tool_bar_new_project.connect('clicked', self.__on_new_project)
        self.tool_bar_new_project.set_tooltip_text('New or Open a project')
        self.__ena(self.tool_bar_new_project)
        tool_bar_hbox.pack_start(self.tool_bar_new_project, expand=True)

        self.tool_bar_close_project = gtk.Button('_Close')
        self.tool_bar_close_project.set_image(stock(gtk.STOCK_CLOSE))
        self.tool_bar_close_project.connect('clicked', self.__on_close_project)
        self.tool_bar_close_project.set_tooltip_text('close current project')
        self.__dis(self.tool_bar_close_project)
        tool_bar_hbox.pack_start(self.tool_bar_close_project, expand=True)

        self.tool_bar_save_project = gtk.Button('_Save')
        self.tool_bar_save_project.set_image(stock(gtk.STOCK_SAVE))
        self.tool_bar_save_project.connect('clicked', self.__on_save_project)
        self.tool_bar_save_project.set_tooltip_text('save change to items and blobs')
        self.__dis(self.tool_bar_save_project)
        tool_bar_hbox.pack_start(self.tool_bar_save_project, expand=True)

        tool_bar_hbox.pack_start(gtk.Label('  '), expand=False)

        self.tool_bar_backup_project = gtk.Button('_Backup')
        self.tool_bar_backup_project.set_image(stock(gtk.STOCK_EXECUTE))
        self.tool_bar_backup_project.connect('clicked', self.__on_backup_project)
        self.tool_bar_backup_project.set_tooltip_text('Backup the project to a zip')
        self.__dis(self.tool_bar_backup_project)
        tool_bar_hbox.pack_start(self.tool_bar_backup_project, expand=True)

        self.tool_bar_search_project = gtk.Button('S_earch')
        self.tool_bar_search_project.set_image(stock(gtk.STOCK_FIND))
        self.tool_bar_search_project.connect('clicked', self.__on_search_project)
        self.tool_bar_search_project.set_tooltip_text('Search files in project')
        self.__dis(self.tool_bar_search_project)
        tool_bar_hbox.pack_start(self.tool_bar_search_project, expand=True)

        self.tool_bar_check_project = gtk.Button('Chec_k')
        self.tool_bar_check_project.set_image(stock(gtk.STOCK_INDEX))
        self.tool_bar_check_project.connect('clicked', self.__on_check_project)
        self.tool_bar_check_project.set_tooltip_text('Check if files in project exists')
        self.__dis(self.tool_bar_check_project)
        tool_bar_hbox.pack_start(self.tool_bar_check_project, expand=True)

        tool_bar_hbox.pack_start(gtk.Label('  '), expand=False)

        self.tool_bar_add_file = gtk.Button('Add _Files')
        self.tool_bar_add_file.set_image(stock(gtk.STOCK_FILE))
        self.tool_bar_add_file.connect('clicked', self.__on_add_files)
        self.tool_bar_add_file.set_tooltip_text('Add files to the project')
        self.__dis(self.tool_bar_add_file)
        tool_bar_hbox.pack_start(self.tool_bar_add_file, expand=True)

        self.tool_bar_add_directory = gtk.Button('Add _Directory')
        self.tool_bar_add_directory.set_image(stock(gtk.STOCK_DIRECTORY))
        self.tool_bar_add_directory.connect('clicked', self.__on_add_directory)
        self.tool_bar_add_directory.set_tooltip_text('Add directory to the project')
        self.__dis(self.tool_bar_add_directory)
        tool_bar_hbox.pack_start(self.tool_bar_add_directory, expand=True)

        tool_bar_hbox.pack_start(gtk.Label('  '), expand=False)

        self.tool_bar_expand_all = gtk.Button('Expand')
        self.tool_bar_expand_all.set_image(stock(gtk.STOCK_FLOPPY))
        self.tool_bar_expand_all.connect('clicked', self.__on_expand_all)
        self.tool_bar_expand_all.set_tooltip_text('expand all')
        self.__dis(self.tool_bar_expand_all)
        #tool_bar_hbox.pack_start(self.tool_bar_expand_all, expand=True)

        self.tool_bar_collapse_all = gtk.Button('Collapse')
        self.tool_bar_collapse_all.connect('clicked', self.__on_collapse_all)
        self.tool_bar_collapse_all.set_tooltip_text('collapse all')
        self.__dis(self.tool_bar_collapse_all)
        #tool_bar_hbox.pack_start(self.tool_bar_collapse_all, expand=True)

        tool_bar_hbox.pack_start(gtk.Label('  '), expand=False)

        self.tool_bar_load_cache = gtk.Button('_Load cache')
        self.tool_bar_load_cache.set_image(stock(gtk.STOCK_CONNECT))
        self.tool_bar_load_cache.connect('clicked', self.__on_load_cache)
        self.tool_bar_load_cache.set_tooltip_text('Load files in ~/.beifen-cache')
        self.__dis(self.tool_bar_load_cache)
        tool_bar_hbox.pack_start(self.tool_bar_load_cache, expand=True)

        self.__body_vbox.pack_start(tool_bar_hbox, expand=False)

    def add_files(self, files):
        model = self.__sort_model.get_model()
        cur_files = self.__project_content.get_files()

        for f in files:
            if os.path.isdir(f):
                size = -1

                if not f.endswith('/'):
                    f = f + '/'
            else:
                try:
                    size = os.path.getsize(f)
                except:
                    size = -1

            if cur_files.has_key(f):
                continue

            col = [ '', f, os.path.splitext(f)[1], size]
            model.append(col)
            self.__project_content.add_path(f)

            # add the parent directory also
            if os.path.isdir(f):
                f = f[:-1]
            dirname = os.path.dirname(f)
            if dirname != '':
                self.add_files([dirname])

        self.__set_dirty(True)

    def add_directory(self, directory):
        self.add_files([directory])

        try:
            ls = os.listdir(directory)
        except:
            print 'access deny'
        else:
            for l in ls:
                temp = os.path.join(directory, l)
                if (os.path.isdir(temp)):
                    self.add_directory(temp)
                else:
                    self.add_files([temp])

    def del_all(self):
        selection = self.__proj_tree_view.get_selection()
        selection.select_all()

        self.del_selected()

    def del_selected(self):

        def match_text(model, path, iter, data):
            [selection, text] = data
            path = model.get(iter, 1)[0]

            # match dir
            if text.endswith('/'):
                if path.startswith(text):
                    selection.select_iter(iter)
            elif path == text:
                    selection.select_iter(iter)

        selection = self.__proj_tree_view.get_selection()
        sorted_model, selected = selection.get_selected_rows()
        model = sorted_model.get_model()
        iters = [sorted_model.convert_iter_to_child_iter(None, sorted_model.get_iter(path)) for path in selected]

        selection.unselect_all()

        if len(iters) > 0:
            self.__set_dirty(True)

        for iter in iters:
            path = model.get(iter, 1)[0]
            self.__proj_tree_view.get_model().foreach(match_text, [selection, path])

        sorted_model, selected = selection.get_selected_rows()
        model = sorted_model.get_model()
        iters = [sorted_model.convert_iter_to_child_iter(None, sorted_model.get_iter(path)) for path in selected]
        for iter in iters:
            path = model.get(iter, 1)[0]
            self.__project_content.del_path(path)
            model.remove(iter)

    def search(self, text):

        def match_text(model, path, iter, data):
            [selection, text] = data
            path = model.get(iter, 1)[0]
            if path.find(text) >= 0:
                selection.select_iter(iter)

        selection = self.__proj_tree_view.get_selection()
        self.__proj_tree_view.get_model().foreach(match_text, [selection, text])

    def __on_tree_view_key_press(self, treeview, key):
        if key.keyval == gtk.keysyms.Delete:
            self.del_selected()

    def __on_tree_view_button_press(self, treeview, event):
        if event.button == 1:
            if event.type == gtk.gdk._2BUTTON_PRESS:
                pathinf = treeview.get_path_at_pos(int(event.x), int(event.y))
                if pathinf:
                    path, col, x, y = pathinf
                    model = treeview.get_model()
                    fn = model.get_value(model.get_iter(path), 1)
                    if os.path.exists(fn):
                        subprocess.Popen(['gnome-open', fn])
            else:
                selection = self.__proj_tree_view.get_selection()
                selection.unselect_all()

    def __create_proj_list(self):
        scroll_window = gtk.ScrolledWindow()
        scroll_window.set_border_width(4)
        scroll_window.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
        scroll_window.set_shadow_type(gtk.SHADOW_NONE)

        self.__body_vbox.pack_start(scroll_window)


        self.__proj_list_store = gtk.ListStore(str, str, str, int)

        self.__sort_model = gtk.TreeModelSort(self.__proj_list_store)
        self.__sort_model.set_sort_column_id(1, gtk.SORT_ASCENDING)

        self.__proj_tree_view = gtk.TreeView(self.__sort_model)
        scroll_window.add(self.__proj_tree_view)

        self.__proj_tree_view.set_rules_hint(True)

        self.__proj_tree_view.get_selection().set_mode(gtk.SELECTION_MULTIPLE)
        self.__proj_tree_view.connect('key-press-event', self.__on_tree_view_key_press)
        self.__proj_tree_view.connect('button-press-event', self.__on_tree_view_button_press)

        self.__proj_tree_view.column = [None]*4
        self.__proj_tree_view.column[0] = gtk.TreeViewColumn('S')
        self.__proj_tree_view.column[1] = gtk.TreeViewColumn('Path')
        self.__proj_tree_view.column[1].set_resizable(True)
        self.__proj_tree_view.column[1].set_expand(True)
        self.__proj_tree_view.column[2] = gtk.TreeViewColumn('Ext')
        self.__proj_tree_view.column[2].set_resizable(True)
        self.__proj_tree_view.column[3] = gtk.TreeViewColumn('Size')
        self.__proj_tree_view.column[3].set_resizable(True)
        self.__proj_tree_view.cell = [None]*4
        for i in range(4):
            self.__proj_tree_view.cell[i] = gtk.CellRendererText()
            self.__proj_tree_view.append_column(self.__proj_tree_view.column[i])
            self.__proj_tree_view.column[i].set_sort_column_id(i)
            self.__proj_tree_view.column[i].pack_start(self.__proj_tree_view.cell[i], True)
            self.__proj_tree_view.column[i].set_attributes(self.__proj_tree_view.cell[i], text=i)


    def __create_status_bar(self):
        self.__status_bar = gtk.Statusbar()
        self.__status_bar_content_id = self.__status_bar.get_context_id('messages')

        #self.__body_vbox.pack_start(self.__status_bar, expand=False)

    def __set_status_text(self, text):
        print text
        self.__status_bar.push(self.__status_bar_content_id, text)

class SearchWindow(gtk.Dialog):
    def get_text(self):
        return self.search_entry.get_text()

    def __init__(self):
        gtk.Window.__init__(self)
        self.set_title('Bei the shit')

        self.__body_vbox = gtk.VBox()
        self.add(self.__body_vbox)

        frame = gtk.Frame('Searching what?')
        self.search_entry = gtk.Entry()
        frame.add(self.search_entry)

        self.vbox.pack_start(frame, True, True, 0)

        self.add_buttons(gtk.STOCK_FIND, gtk.RESPONSE_OK, gtk.STOCK_CANCEL, gtk.RESPONSE_CANCEL)

        self.show_all()

if __name__ == '__main__':
    gobject.threads_init()
    MainWindow()
    gtk.main()
